package com.star.springai.chat.controller;

import java.util.List;
import java.util.Map;

import lombok.extern.slf4j.Slf4j;

import com.star.springai.chat.service.MultipartFileResource;
import com.star.springai.chat.service.MyPagePdfDocumentReader;
import com.star.springai.chat.service.MyTikaDocumentReader;
import com.star.springai.chat.service.MyTokenTextSplitter;

import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

@RestController
@RequestMapping("/vectorStore")
@Slf4j
public class VectorStoreController {
    @Autowired
    private VectorStore vectorStore;
    @Autowired
    private MyPagePdfDocumentReader pdfReader;
    @Autowired
    private MyTikaDocumentReader tikaDocumentReader;
    @Autowired
    private MyTokenTextSplitter tokenTextSplitter;



    public record UserKnowledge(String knowledgeType, String knowledgeContent) {
    }

    /**
     * 用户传入信息存入知识库
     */
    @PostMapping("/addKnowledge")
    public Boolean addDocument(@RequestBody UserKnowledge userKnowledge) {

        Document document = new Document(
                userKnowledge.knowledgeContent,
                Map.of(userKnowledge.knowledgeType, userKnowledge.knowledgeContent));
        vectorStore.add(List.of(document));
        return true;
    }

    /**
     * 上传PDF文件并存入知识库
     */
    @PostMapping("/uploadPdf")
    public List<Document> uploadPdf(@RequestParam("file") MultipartFile file) {
        try {
            // 将MultipartFile包装成Resource
            MultipartFileResource resource = new MultipartFileResource(file);

            // 读取PDF文档
            List<Document> documents = pdfReader.getDocsFromPdf(resource);

            // 存入向量数据库
            vectorStore.add(documents);

            return documents;
        } catch (Exception e) {
            log.error("处理PDF文件时发生错误", e);
            throw new RuntimeException("文件处理失败: " + e.getMessage());
        }
    }

    /**
     * 上传其他文档文件并存入知识库
     */
    @PostMapping("/uploadDocument")
    public List<Document> uploadDocument(@RequestParam("file") MultipartFile file) {
        try {
            // 将MultipartFile包装成Resource
            MultipartFileResource resource = new MultipartFileResource(file);

            // 读取文档
            List<Document> documents = tikaDocumentReader.loadText(resource);

            // 存入向量数据库
            vectorStore.add(documents);

            return documents;
        } catch (Exception e) {
            log.error("处理文档文件时发生错误", e);
            throw new RuntimeException("文件处理失败: " + e.getMessage());
        }
    }

    /**
     * 上传PDF文件并拆分后存入知识库
     */
    @PostMapping("/uploadPdfAndSplit")
    public List<Document> uploadPdfAndSplit(@RequestParam("file") MultipartFile file) {
        try {
            // 将MultipartFile包装成Resource
            MultipartFileResource resource = new MultipartFileResource(file);

            // 读取PDF文档
            List<Document> documents = pdfReader.getDocsFromPdf(resource);

            // 拆分文档
            List<Document> documentsResult = tokenTextSplitter.splitCustomized(documents);

            // 存入向量数据库
            vectorStore.add(documentsResult);

            return documentsResult;
        } catch (Exception e) {
            log.error("处理PDF文件时发生错误", e);
            throw new RuntimeException("文件处理失败: " + e.getMessage());
        }
    }


    @GetMapping("/getFromVector")
    public List<Document> getFromVector(String knowledge) {
        // Retrieve documents similar to a query
        return this.vectorStore.similaritySearch(SearchRequest.builder().query(knowledge)
                .topK(5)
                .similarityThreshold(0.8)
                .build());

    }


}
